from .Vector import Vector
from .Matrix import Matrix
from ._global import is_zero

class LinearSystemExt:
    def __init__(self, A, b = None):
        assert b is None or A.row_num() == len(b), "A row must == b len"
        self._m = A.row_num()
        self._n = A.col_num()

        if b is None:
            self.Ab = [ A.row_vector(i) for i in range(self._m)]

        if isinstance(b, Vector):
            self.Ab = [ Vector(A.row_vector(i).underlying_list() + [b[i]])
                        for i in range(self._m)]

        if isinstance(b, Matrix):
            self.Ab = [ Vector(A.row_vector(i).underlying_list() + b.row_vector(i).underlying_list())
                        for i in range(self._m)]

        self.pivots = []

    def gauss_jordan_elimination(self):
        self._forward()
        self._backward()

        for i in range(len(self.pivots), self._m):
            if not is_zero(self.Ab[i][-1]):
                return False
        return True 

    def _forward(self):
        i, k = 0, 0
        while i < self._m and k < self._n:
            # 交换最大行，提高数值计算的精度
            max_row = self._max_row(i, k, self._m)
            self.Ab[i], self.Ab[max_row] = self.Ab[max_row], self.Ab[i]

            if is_zero(self.Ab[i][k]):
                k += 1
            else:
                #将主元归一，Ab[i][k]为主元
                self.Ab[i] = self.Ab[i] / self.Ab[i][k]

                #同一列下方元素置0
                for j in range(i + 1, self._m):
                    self.Ab[j] = self.Ab[j] - self.Ab[j][k] * self.Ab[i]
                self.pivots.append(k)
                i += 1

    def _backward(self):
        n = len(self.pivots)
        for i in range(n - 1, -1, -1):
            k = self.pivots[i]
            # Ab[i][k]为主元
            for j in range(i - 1, -1, -1):
                self.Ab[j] = self.Ab[j] - self.Ab[i] * self.Ab[j][k]

    def _max_row(self, idx_i, idx_j, n):
        best, ret = self.Ab[idx_i][idx_j], idx_i
        for i in range(idx_i + 1, n):
            if self.Ab[i][idx_j] > best:
                best, ret = self.Ab[i][idx_j], i
        return ret

    def fancy_print(self):
        for i in range(self._m):
            print(" ".join(str(self.Ab[i][j]) for j in range(self._n)), end = " ")
            print("|", self.Ab[i][-1])


def inv(A):
    """返回A的逆矩阵"""
    if A.row_num() != A.col_num():
        return None

    n = A.row_num()
    ls = LinearSystemExt(A, Matrix.identity(n))

    if not ls.gauss_jordan_elimination():
        return None

    invA = [[row[i] for i in range(n, 2 *n)] for row in ls.Ab]

    return Matrix(invA) 

def rank(A):
    ls = LinearSystemExt(A)
    ls.gauss_jordan_elimination()

    zero = Vector.zero(A.col_num())
    return sum([row != zero for row in ls.Ab])